import os
import sys
import io
import base64
from flask import Flask, request, jsonify, render_template
import numpy as np
import cv2
import torch
from PIL import Image

# 关键功能：将父目录加入模块搜索路径，便于导入模型与配置
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..')))
# 关键功能：再加入上一级的父目录，使得 "Child_Adenoid.models" 绝对导入在模型内部也可用
sys.path.append(os.path.abspath(os.path.join(os.path.dirname(__file__), '..', '..')))
from models.AbgNet1 import AbgNet1
from config import Config

app = Flask(__name__, template_folder=os.path.join(os.path.dirname(__file__), 'templates'))
try:
    from flask_cors import CORS
    CORS(app, resources={r"/*": {"origins": "*"}})
except Exception:
    # 关键功能：CORS 兼容处理（无 flask_cors 时的降级方案）
    @app.after_request
    def add_cors_headers(resp):
        resp.headers['Access-Control-Allow-Origin'] = '*'
        resp.headers['Access-Control-Allow-Headers'] = 'Content-Type, Authorization, X-Requested-With, Accept, Origin'
        resp.headers['Access-Control-Allow-Methods'] = 'GET, POST, OPTIONS'
        return resp

    @app.before_request
    def handle_options_preflight():
        if request.method == 'OPTIONS':
            resp = app.make_response('')
            resp.status_code = 204
            return resp

# 关键功能：全局加载模型，加速后续请求推理
device = Config.DEVICE
MODEL_PATH = os.path.abspath(os.path.join(os.path.dirname(__file__), '..', 'saved_models', 'best_model', 'best_model_fold4.pth'))
model = AbgNet1(num_classes=2).to(device)
try:
    state_dict = torch.load(MODEL_PATH, map_location=device)
    # 兼容 DataParallel 保存的权重
    if isinstance(state_dict, dict) and 'module.' in next(iter(state_dict.keys())):
        from collections import OrderedDict
        new_state = OrderedDict()
        for k, v in state_dict.items():
            new_key = k[7:] if k.startswith('module.') else k
            new_state[new_key] = v
        state_dict = new_state
    model.load_state_dict(state_dict)
    model.eval()
    print(f"[Flask] 模型加载成功: {MODEL_PATH}")
except Exception as e:
    print(f"[Flask] 模型加载失败，请检查路径: {e}")

# 后处理依赖
import math


def _fill_holes(mask):
    # 使用 OpenCV 的洪泛填充实现孔洞填充，避免依赖 SciPy/Skimage
    h, w = mask.shape[:2]
    # 创建边界填充的反掩模
    flood = np.zeros((h+2, w+2), dtype=np.uint8)
    # 反相：背景为1，前景为0，便于从边界填充
    inv = (mask == 0).astype(np.uint8)
    # 从边界 (0,0) 进行洪泛填充
    cv2.floodFill(inv, flood, (0, 0), 2)
    # 将未被填充的背景视为孔洞（值为0），与原掩模合并
    holes = (inv == 0).astype(np.uint8)
    filled = np.clip(mask.astype(np.uint8) | holes, 0, 1)
    return filled


def _remove_small_objects(mask, min_area=100):
    # 使用连通域统计移除小尺寸区域
    num, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), connectivity=8)
    cleaned = np.zeros_like(mask, dtype=np.uint8)
    for i in range(1, num):  # 0 是背景
        area = stats[i, cv2.CC_STAT_AREA]
        if area >= min_area:
            cleaned[labels == i] = 1
    return cleaned


def _keep_largest(mask):
    # 保留最大连通域
    num, labels, stats, _ = cv2.connectedComponentsWithStats(mask.astype(np.uint8), connectivity=8)
    if num <= 1:
        return mask
    # 找到面积最大（除背景）的标签
    areas = stats[1:, cv2.CC_STAT_AREA]
    max_idx = 1 + int(np.argmax(areas))
    kept = (labels == max_idx).astype(np.uint8)
    return kept


# 关键功能：模型输出的分割结果后处理为干净的二值掩模
def postprocess_prediction(pred):
    # 关键功能：纯 OpenCV/NumPy 后处理，避免 SciPy/Skimage 依赖
    pred_prob = torch.sigmoid(pred).cpu().numpy()[0]
    adenoid_prob = cv2.GaussianBlur(pred_prob[0], (0, 0), sigmaX=1.1, sigmaY=1.1)
    cavity_prob = cv2.GaussianBlur(pred_prob[1], (0, 0), sigmaX=1.1, sigmaY=1.1)

    adenoid_pred = (adenoid_prob > 0.5).astype(np.uint8)
    cavity_pred = (cavity_prob > 0.5).astype(np.uint8)

    adenoid_pred = _fill_holes(adenoid_pred)
    cavity_pred = _fill_holes(cavity_pred)
    adenoid_pred = _remove_small_objects(adenoid_pred)
    cavity_pred = _remove_small_objects(cavity_pred)
    adenoid_pred = _keep_largest(adenoid_pred)
    cavity_pred = _keep_largest(cavity_pred)

    kernel = cv2.getStructuringElement(cv2.MORPH_ELLIPSE, (3, 3))
    adenoid_pred = cv2.morphologyEx(adenoid_pred, cv2.MORPH_CLOSE, kernel)
    cavity_pred = cv2.morphologyEx(cavity_pred, cv2.MORPH_CLOSE, kernel)

    return adenoid_pred, cavity_pred


def generate_contour_from_mask(mask):
    kernel = np.ones((3, 3), np.uint8)
    dilated = cv2.dilate(mask, kernel, iterations=1)
    eroded = cv2.erode(mask, kernel, iterations=1)
    contour = dilated - eroded
    return np.clip(contour, 0, 1)


def calculate_adenoid_ratio(adenoid_pred, cavity_pred):
    ad = (adenoid_pred > 0.5).astype(np.uint8)
    cv = (cavity_pred > 0.5).astype(np.uint8)
    red = np.sum(ad)
    total = np.sum(np.logical_or(ad, cv))
    ratio = red / total if total > 0 else 0.0
    return ratio


@app.route('/', methods=['GET'])
def index():
    return render_template('index.html')


# 关键功能：Flask 上传接口，完成推理并返回可视化结果
@app.route('/predict', methods=['POST', 'OPTIONS'])
def predict():
    # 关键功能：处理跨域预检请求
    if request.method == 'OPTIONS':
        return '', 204
    if 'file' not in request.files:
        return jsonify({'error': '未找到文件字段 file'}), 400

    file = request.files['file']
    try:
        pil_image = Image.open(file.stream).convert('RGB')
    except Exception:
        return jsonify({'error': '无效图片文件'}), 400

    original_img = np.array(pil_image)
    orig_h, orig_w = original_img.shape[:2]

    # 预处理：Resize 与归一化（与训练保持一致）
    img_resized = cv2.resize(original_img, Config.IMAGE_SIZE)
    img_tensor = torch.from_numpy(img_resized).permute(2, 0, 1).float() / 255.0
    img_tensor = img_tensor.unsqueeze(0).to(device)

    with torch.no_grad():
        pred = model(img_tensor)['final']
        if pred.shape[2:] != Config.IMAGE_SIZE[::-1]:
            pred = torch.nn.functional.interpolate(
                pred, size=Config.IMAGE_SIZE[::-1], mode='bicubic', align_corners=True
            )

    adenoid_pred, cavity_pred = postprocess_prediction(pred)
    ad_mask = cv2.resize(adenoid_pred.astype(np.uint8), (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
    cv_mask = cv2.resize(cavity_pred.astype(np.uint8), (orig_w, orig_h), interpolation=cv2.INTER_NEAREST)
    ratio = calculate_adenoid_ratio(ad_mask, cv_mask)

    ad_contour = generate_contour_from_mask(ad_mask)
    cv_contour = generate_contour_from_mask(cv_mask)
    display_img = original_img.copy()
    display_img[ad_contour > 0] = [255, 0, 0]
    display_img[cv_contour > 0] = [0, 255, 0]

    res_pil = Image.fromarray(display_img)
    buff = io.BytesIO()
    res_pil.save(buff, format='PNG')
    base64_str = base64.b64encode(buff.getvalue()).decode('utf-8')

    return jsonify({
        'image': f'data:image/png;base64,{base64_str}',
        'ratio': float(ratio),
        'diagnosis': '堵塞严重' if ratio >= 0.6 else ('轻微堵塞' if ratio >= 0.3 else '正常')
    })


# 路由别名：对齐前端接口路径
@app.route('/ai-detection/detect', methods=['POST', 'OPTIONS'])
def predict_alias():
    # 关键功能：处理跨域预检请求（别名路由）
    if request.method == 'OPTIONS':
        return '', 204
    return predict()


if __name__ == '__main__':
    # Flask 原生开发服务器，适合示范用途
    app.run(host='127.0.0.1', port=8001, debug=True)

